//
// Created by hjhu on 2024/7/2.
//

#ifndef TEST_WAV2LIP_CPP_HOURGLASS_H
#define TEST_WAV2LIP_CPP_HOURGLASS_H

#include <torch/torch.h>
#include "BaseModule.h"
namespace coastal {

    using namespace torch;
    using namespace std;

    class HourGlassImpl : public coastal::BaseModuleImpl {

    public:

        HourGlassImpl(int64_t num_modules, int64_t depth, int64_t num_features);

        torch::Tensor forward(torch::Tensor x);

    private:
        int64_t num_modules;
        int64_t depth;
        int64_t features;

        void generate_network(int64_t level);

        torch::Tensor _forward(int64_t level, torch::Tensor inp);



    };

} // coastal

#endif //TEST_WAV2LIP_CPP_HOURGLASS_H
